from mmdet.core import bbox2result
from ..builder import DETECTORS
from .single_stage import SingleStageDetector

@DETECTORS.register_module()
class DETR(SingleStageDetector):
    r"""Implementation of 'DETR:End-to-End Object Detection with
    Transformers <https://arxiv.org/pdf/2005.12872>'_"""

    def __init__(self,
                 backbone,
                 bbox_head,
                 train_cfg=None,
                 test_cfg=None,
                 pretrained=None):
        super(DETR, self).__init__(backbone,None,bbox_head,train_cfg,
                                   test_cfg,pretrained)

    def simple_test(self,img,img_metas,rescale=False):
        """Test function out test time augmentation.

        A:
            imgs(list[torch.Tensor]):List of multiple images
            img_metas(list[dict]):List of image information.
            rescale(bool,optional):Whether to rescale the results.
            Default to false.

        Returns:
            list[list[np.ndarray]]:BBox results of each image and classes.
            The outer list corresponds to each image.The inner list corresponds to each class.
        """
        batch_size = len(img_metas)
        assert batch_size==1,'Currently only batch_size 1 for inference'\
            f'mode is supported.Found batch_size{batch_size}.'
        x = self.extract_feat(img)
        outs = self.bbox_head(x,img_metas)
        bbox_list = self.bbox_head.get_bboxes(
            *outs,img_metas,rescale=rescale
        )

        bbox_results = [
            bbox2result(det_bboxes,det_labels,self.bbox_head.num_classes)
            for det_bboxes,det_labels in bbox_list
        ]

        return bbox_results
